import os
import time
import argparse

import tensorflow as tf

FLAGS = None

def train(fraction=1.0):
    config = tf.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = fraction

    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
    b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
    c = tf.matmul(a, b)
    # Creates a session with log_device_placement set to True.
    config = tf.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = fraction
    sess = tf.Session(config=config)
    
    start=time.time()
    # Runs the op.
    while True:
        sess.run(c)


if __name__ == '__main__':
    # os.environ["CUDA_VISIBLE_DEVICES"] = "1" # 指定使用哪一块显卡，与k8s中的环境变量作用相同
    parser = argparse.ArgumentParser()
    parser.add_argument('--total', type=float, default=1000,
                      help='Total GPU memory.')
    parser.add_argument('--allocated', type=float, default=1000,
                      help='Allocated GPU memory.')
    FLAGS, unparsed = parser.parse_known_args()
    # fraction = FLAGS.allocated / FLAGS.total * 0.85
    fraction = round( FLAGS.allocated * 0.2 / FLAGS.total , 1 )

    print(fraction) # fraction 默认值为0.7，该程序最多使用总资源的70%
    train(fraction)